from argparse import ArgumentParser

from lagent.llms import HFTransformerCasualLM
from lagent.llms.meta_template import INTERNLM2_META as META


def parse_args():
    """
    解析命令行参数。

    Returns:
        argparse.Namespace: 包含解析后的参数
            - path: 模型路径
            - mode: 对话模式（chat 或 generate）
    """
    parser = ArgumentParser(description='chatbot')
    parser.add_argument(
        '--path',
        type=str,
        default='/root/autodl-tmp/models/internlm2_5-1_8b-chat',  # 默认使用 InternLM2-Chat-20B 模型
        help='The path to the model')
    parser.add_argument(
        '--mode',
        type=str,
        default='chat',  # 默认使用对话模式
        help='Completion through chat or generate')
    args = parser.parse_args()
    return args


def main():
    """
    主函数：实现一个简单的命令行聊天机器人。
    支持多行输入，可以通过 'exit' 命令退出。
    """
    args = parse_args()
    # 初始化 HuggingFace Transformer 模型
    model = HFTransformerCasualLM(
        path=args.path,  # 模型路径
        meta_template=META,  # 使用 InternLM2 的元模板
        max_new_tokens=1024,  # 最大生成令牌数
        stop_words=['<|im_end|>'])  # 停止词，用于终止生成

    # 设置生成参数
    model.gen_params.update(
        do_sample=True,  # 启用采样模式
        top_p=0.8,  # 采样时的累积概率阈值
        top_k=None,  # 不限制候选令牌数量
        temperature=0.1,  # 温度参数，控制生成的随机性
        repetition_penalty=1.0,  # 重复惩罚系数
        batch_size=1,  # 添加批次大小参数
    )

    def input_prompt():
        """获取用户输入，支持多行输入"""
        print('\ndouble enter to end input >>> ', end='', flush=True)
        sentinel = ''
        return '\n'.join(iter(input, sentinel))

    # 开始对话循环
    history = []  # 存储对话历史
    while True:
        try:
            prompt = input_prompt()
        except UnicodeDecodeError:
            print('UnicodeDecodeError')
            continue
        
        if prompt == 'exit':
            exit(0)
            
        # 将用户输入添加到历史记录
        history.append(dict(role='user', content=prompt))
        
        # 如果是生成模式，则不保留历史记录
        if args.mode == 'generate':
            history = [dict(role='user', content=prompt)]
            
        # 使用 stream_chat 方法
        print('\nInternLm2：', end='')
        current_length = 0  # 记录已输出的长度
        try:
            for status, response, _ in model.stream_chat(history):
                # 只打印新生成的内容
                print(response[current_length:], end='', flush=True)
                current_length = len(response)
                
            # 将模型响应添加到历史记录
            history.append(dict(role='assistant', content=response))
            print('')  # 换行
        except Exception as e:
            print(f"\nError: {str(e)}")
            continue


if __name__ == '__main__':
    main()
